{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FIGConvNet inference and visualization notebook\n",
    "\n",
    "This notebook demonstrates how to use pre-trained FIGConvNet model to perform inference\n",
    "on DrivAerNet dataset.\n",
    "\n",
    "The following items are required and need to be downloaded separately:\n",
    "\n",
    "* Pre-trained [FIGConvNet checkpoint](to_be_provided).\n",
    "* [DrivAerNet](https://github.com/Mohamedelrefaie/DrivAerNet/tree/main/DrivAerNet_v1) dataset.\n",
    "    The dataset needs to be converted to a Webdataset format. For simplicity, the small subset\n",
    "    of the dataset has been already converted to Webdataset format and can be used in this\n",
    "    example as-is.\n",
    "\n",
    "The inputs to the model are:\n",
    "* Point cloud representing the surface of the car.\n",
    "\n",
    "The outputs of the model are:\n",
    "* Pressure at each surface point.\n",
    "* Wall shear stresses at each surface point.\n",
    "\n",
    "Before we begin, let's import some common packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warp 1.0.2 initialized:\n",
      "   CUDA Toolkit 11.5, Driver 12.2\n",
      "   Devices:\n",
      "     \"cpu\"      : \"x86_64\"\n",
      "     \"cuda:0\"   : \"NVIDIA GeForce RTX 3090\" (24 GiB, sm_86, mempool enabled)\n",
      "     \"cuda:1\"   : \"NVIDIA TITAN RTX\" (24 GiB, sm_75, mempool enabled)\n",
      "   CUDA peer access:\n",
      "     Not supported\n",
      "   Kernel cache:\n",
      "     /home/du/.cache/warp/1.0.2\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import pyvista as pv\n",
    "import torch\n",
    "import warp as wp\n",
    "\n",
    "if sys.path[0] != \"..\":\n",
    "    sys.path.insert(0, \"..\")\n",
    "\n",
    "device = torch.device(\"cuda:0\")\n",
    "torch.cuda.device(device)\n",
    "wp.init()\n",
    "wp.set_device(str(device))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset visualization\n",
    "\n",
    "This section provides visualizations of the original, mesh-based dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to the dataset and pressure VTK files.\n",
    "# Note: update `drivaer_orig_path` as needed.\n",
    "drivaer_orig_path = Path(\"/data/src/physicsnemo/data/DrivAerNet/mini\")\n",
    "vtk_path = drivaer_orig_path / \"SurfacePressureVTK\"\n",
    "\n",
    "output_dir = drivaer_orig_path.parent / \"vis\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e670e1bc4d24aae87943a08fe91197d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Widget(value='<iframe src=\"http://localhost:42125/index.html?ui=P_0x7f21d748c2b0_0&reconnect=auto\" class=\"pyvi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Design id 0001 corresponds to DrivAer_F_D_WM_WW_0001.vtk file which belongs to the test set.\n",
    "design_ids = [\"0001\"]\n",
    "\n",
    "\n",
    "camera_position = [\n",
    "    (-3.0, -4.5, 5),\n",
    "    (1, 0, 0.6),\n",
    "    (0.1, 0.1, 0.9),\n",
    "]\n",
    "\n",
    "\n",
    "def create_mesh_vis(\n",
    "    mesh: pv.PolyData,\n",
    "    plotter: pv.Plotter,\n",
    "    camera_position,\n",
    "):\n",
    "    plotter.subplot(0, 0)\n",
    "    # Solid mesh visualization\n",
    "    plotter.add_mesh(mesh, color=\"lightgrey\")\n",
    "    plotter.camera_position = camera_position\n",
    "    plotter.add_text(\"Mesh\", position=\"upper_left\")\n",
    "\n",
    "\n",
    "def create_mesh_vis_gt(\n",
    "    mesh: pv.PolyData,\n",
    "    scalar_name: str,\n",
    "    plotter: pv.Plotter,\n",
    "    camera_position,\n",
    "):\n",
    "    plotter.subplot(0, 1)\n",
    "    # Solid mesh visualization with scalar.\n",
    "    plotter.add_mesh(\n",
    "        mesh,\n",
    "        scalars=scalar_name,\n",
    "        cmap=\"jet\",\n",
    "        clim=(-600, 400),\n",
    "        show_scalar_bar=False,\n",
    "    )\n",
    "    plotter.camera_position = camera_position\n",
    "    plotter.add_text(\"GT Pressure\", position=\"upper_right\")\n",
    "\n",
    "\n",
    "def visualize_meshes(\n",
    "    output_dir: Path = None,\n",
    "):\n",
    "    for design in design_ids:\n",
    "        mesh_file = vtk_path / f\"DrivAer_F_D_WM_WW_{design}.vtk\"\n",
    "        mesh = pv.read(mesh_file)\n",
    "\n",
    "        plotter = pv.Plotter(shape=(1, 2))\n",
    "        # Create input mesh vis.\n",
    "        create_mesh_vis(mesh, plotter, camera_position)\n",
    "        # Create GT pressure vis.\n",
    "        create_mesh_vis_gt(mesh, \"p\", plotter, camera_position)\n",
    "\n",
    "        if output_dir:\n",
    "            output_dir.mkdir(parents=True, exist_ok=True)\n",
    "            plotter.save_graphic(output_dir / f\"{design}_gt_p_mesh.pdf\")\n",
    "        else:\n",
    "            plotter.show()\n",
    "\n",
    "\n",
    "output_dir = None\n",
    "# Uncomment and update the below to render to pdf files instead.\n",
    "# output_dir = Path(\"/data/src/physicsnemo/data/DrivAerNet/vis\")\n",
    "visualize_meshes(output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running model inference\n",
    "\n",
    "This section provides an example of running model inference on the examples from DrivAerNet dataset.\n",
    "\n",
    "For simplicity, the same dataloader is used in inference as in training (Webdataset-based)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the dataloader\n",
    "\n",
    "Instantiate the dataloader first. Please update the path to the Webdataset as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import src.data\n",
    "\n",
    "\n",
    "num_points = 65536  # 2048\n",
    "dataset = src.data.DrivAerNetDataModule(\n",
    "    drivaer_orig_path.parent / \"drivaernet_webdataset\",\n",
    "    num_points=num_points,\n",
    "    preprocessors=[src.data.drivaernet_datamodule.DrivAerNetPreprocessor(num_points)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the model\n",
    "\n",
    "The following code creates a model and loads pre-trained weights from a checkpoint file.\n",
    "The model must be instantiated with exactly the same arguments that were used\n",
    "when the model was trained.\n",
    "\n",
    "Note: update the path to the checkpoint as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/torch/functional.py:512: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/TensorShape.cpp:3559.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FIGConvUNetDrivAerNet(\n",
       "  (point_feature_to_grids): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): PointFeatureToGrid(\n",
       "        (conv): PointFeatureConv(in_channels=16 out_channels=16 search_type=radius reductions=['mean'] rel_pos_encode=True)\n",
       "      )\n",
       "      (1): GridFeatureMemoryFormatConverter(memory_format=GridFeaturesMemoryFormat.b_xc_y_z)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): PointFeatureToGrid(\n",
       "        (conv): PointFeatureConv(in_channels=16 out_channels=16 search_type=radius reductions=['mean'] rel_pos_encode=True)\n",
       "      )\n",
       "      (1): GridFeatureMemoryFormatConverter(memory_format=GridFeaturesMemoryFormat.b_yc_x_z)\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): PointFeatureToGrid(\n",
       "        (conv): PointFeatureConv(in_channels=16 out_channels=16 search_type=radius reductions=['mean'] rel_pos_encode=True)\n",
       "      )\n",
       "      (1): GridFeatureMemoryFormatConverter(memory_format=GridFeaturesMemoryFormat.b_zc_x_y)\n",
       "    )\n",
       "  )\n",
       "  (down_blocks): ModuleList(\n",
       "    (0-1): 2 x Sequential(\n",
       "      (0): GridFeatureConv2DBlocksAndIntraCommunication(\n",
       "        (convs): ModuleList(\n",
       "          (0): GridFeatureConv2dBlock(\n",
       "            (conv1): GridFeatureConv2d(\n",
       "              (conv): Conv2d(80, 80, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2))\n",
       "            )\n",
       "            (conv2): GridFeatureConv2d(\n",
       "              (conv): Conv2d(80, 80, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "            )\n",
       "            (norm1): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((80,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (norm2): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((80,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (shortcut): GridFeatureConv2d(\n",
       "              (conv): Conv2d(80, 80, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (pad_to_match): GridFeaturePadToMatch()\n",
       "            (nonlinear): GridFeatureTransform(\n",
       "              (feature_transform): GELU(approximate='none')\n",
       "            )\n",
       "          )\n",
       "          (1): GridFeatureConv2dBlock(\n",
       "            (conv1): GridFeatureConv2d(\n",
       "              (conv): Conv2d(48, 48, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2))\n",
       "            )\n",
       "            (conv2): GridFeatureConv2d(\n",
       "              (conv): Conv2d(48, 48, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "            )\n",
       "            (norm1): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((48,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (norm2): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((48,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (shortcut): GridFeatureConv2d(\n",
       "              (conv): Conv2d(48, 48, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (pad_to_match): GridFeaturePadToMatch()\n",
       "            (nonlinear): GridFeatureTransform(\n",
       "              (feature_transform): GELU(approximate='none')\n",
       "            )\n",
       "          )\n",
       "          (2): GridFeatureConv2dBlock(\n",
       "            (conv1): GridFeatureConv2d(\n",
       "              (conv): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2))\n",
       "            )\n",
       "            (conv2): GridFeatureConv2d(\n",
       "              (conv): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "            )\n",
       "            (norm1): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((32,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (norm2): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((32,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (shortcut): GridFeatureConv2d(\n",
       "              (conv): Conv2d(32, 32, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (pad_to_match): GridFeaturePadToMatch()\n",
       "            (nonlinear): GridFeatureTransform(\n",
       "              (feature_transform): GELU(approximate='none')\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (intra_communications): GridFeatureGroupIntraCommunications(\n",
       "          (intra_communications): ModuleList(\n",
       "            (0): GridFeaturesGroupIntraCommunication()\n",
       "          )\n",
       "          (grid_cat): GridFeatureGroupCat(\n",
       "            (grid_cat): GridFeatureCat()\n",
       "          )\n",
       "        )\n",
       "        (proj): Identity()\n",
       "        (nonlinear): GridFeatureGroupTransform(\n",
       "          (transform): GELU(approximate='none')\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (up_blocks): ModuleList(\n",
       "    (0-1): 2 x Sequential(\n",
       "      (0): GridFeatureConv2DBlocksAndIntraCommunication(\n",
       "        (convs): ModuleList(\n",
       "          (0): GridFeatureConv2dBlock(\n",
       "            (conv1): GridFeatureConv2d(\n",
       "              (conv): ConvTranspose2d(80, 80, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (conv2): GridFeatureConv2d(\n",
       "              (conv): Conv2d(80, 80, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "            )\n",
       "            (norm1): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((80,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (norm2): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((80,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (shortcut): GridFeatureConv2d(\n",
       "              (conv): ConvTranspose2d(80, 80, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (pad_to_match): GridFeaturePadToMatch()\n",
       "            (nonlinear): GridFeatureTransform(\n",
       "              (feature_transform): GELU(approximate='none')\n",
       "            )\n",
       "          )\n",
       "          (1): GridFeatureConv2dBlock(\n",
       "            (conv1): GridFeatureConv2d(\n",
       "              (conv): ConvTranspose2d(48, 48, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (conv2): GridFeatureConv2d(\n",
       "              (conv): Conv2d(48, 48, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "            )\n",
       "            (norm1): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((48,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (norm2): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((48,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (shortcut): GridFeatureConv2d(\n",
       "              (conv): ConvTranspose2d(48, 48, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (pad_to_match): GridFeaturePadToMatch()\n",
       "            (nonlinear): GridFeatureTransform(\n",
       "              (feature_transform): GELU(approximate='none')\n",
       "            )\n",
       "          )\n",
       "          (2): GridFeatureConv2dBlock(\n",
       "            (conv1): GridFeatureConv2d(\n",
       "              (conv): ConvTranspose2d(32, 32, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (conv2): GridFeatureConv2d(\n",
       "              (conv): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "            )\n",
       "            (norm1): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((32,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (norm2): GridFeatureTransform(\n",
       "              (feature_transform): LayerNorm2d((32,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "            (shortcut): GridFeatureConv2d(\n",
       "              (conv): ConvTranspose2d(32, 32, kernel_size=(2, 2), stride=(2, 2))\n",
       "            )\n",
       "            (pad_to_match): GridFeaturePadToMatch()\n",
       "            (nonlinear): GridFeatureTransform(\n",
       "              (feature_transform): GELU(approximate='none')\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (intra_communications): GridFeatureGroupIntraCommunications(\n",
       "          (intra_communications): ModuleList(\n",
       "            (0): GridFeaturesGroupIntraCommunication()\n",
       "          )\n",
       "          (grid_cat): GridFeatureGroupCat(\n",
       "            (grid_cat): GridFeatureCat()\n",
       "          )\n",
       "        )\n",
       "        (proj): Identity()\n",
       "        (nonlinear): GridFeatureGroupTransform(\n",
       "          (transform): GELU(approximate='none')\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (convert_to_orig): GridFeatureMemoryFormatConverter(memory_format=GridFeaturesMemoryFormat.b_x_y_z_c)\n",
       "  (grid_pools): ModuleList(\n",
       "    (0): GridFeatureGroupPool(\n",
       "      (pools): ModuleList(\n",
       "        (0): GridFeaturePool(\n",
       "          (conv): Conv2d(80, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (pool): AdaptiveMaxPool1d(output_size=1)\n",
       "          (norm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (1): GridFeaturePool(\n",
       "          (conv): Conv2d(48, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (pool): AdaptiveMaxPool1d(output_size=1)\n",
       "          (norm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (2): GridFeaturePool(\n",
       "          (conv): Conv2d(32, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (pool): AdaptiveMaxPool1d(output_size=1)\n",
       "          (norm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (mlp): MLP(\n",
       "    (layers): ModuleList(\n",
       "      (0): ResidualLinearBlock(\n",
       "        (blocks): Sequential(\n",
       "          (0): Linear(in_features=6144, out_features=6144, bias=True)\n",
       "          (1): LayerNorm((6144,), eps=1e-05, elementwise_affine=True)\n",
       "          (2): GELU(approximate='none')\n",
       "          (3): Linear(in_features=6144, out_features=2048, bias=True)\n",
       "          (4): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (shortcut): Linear(in_features=6144, out_features=2048, bias=True)\n",
       "        (activation): GELU(approximate='none')\n",
       "      )\n",
       "      (1): ResidualLinearBlock(\n",
       "        (blocks): Sequential(\n",
       "          (0): Linear(in_features=2048, out_features=2048, bias=True)\n",
       "          (1): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "          (2): GELU(approximate='none')\n",
       "          (3): Linear(in_features=2048, out_features=2048, bias=True)\n",
       "          (4): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (shortcut): Identity()\n",
       "        (activation): GELU(approximate='none')\n",
       "      )\n",
       "      (2): LinearBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Linear(in_features=2048, out_features=2048, bias=False)\n",
       "          (1): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "          (2): GELU(approximate='none')\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (mlp_projection): Linear(in_features=2048, out_features=1, bias=True)\n",
       "  (to_point): GridFeatureGroupToPoint(\n",
       "    (conv_list): ModuleList(\n",
       "      (0-2): 3 x GridFeatureToPoint(\n",
       "        (conv): GridFeatureToPointGraphConv(\n",
       "          (conv): PointFeatureConv(in_channels=16 out_channels=16 search_type=radius reductions=['mean'] rel_pos_encode=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (projection): PointFeatureTransform(\n",
       "    (feature_transform): Sequential(\n",
       "      (0): Linear(in_features=32, out_features=32, bias=True)\n",
       "      (1): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
       "      (2): GELU(approximate='none')\n",
       "      (3): Linear(in_features=32, out_features=1, bias=True)\n",
       "    )\n",
       "  )\n",
       "  (pad_to_match): GridFeatureGroupPadToMatch(\n",
       "    (match): GridFeaturePadToMatch()\n",
       "  )\n",
       "  (vertex_to_point_features): VerticesToPointFeatures(\n",
       "    (pos_embed): SinusoidalEncoding()\n",
       "    (mlp): MLP(\n",
       "      (layers): ModuleList(\n",
       "        (0): LinearBlock(\n",
       "          (block): Sequential(\n",
       "            (0): Linear(in_features=96, out_features=16, bias=False)\n",
       "            (1): LayerNorm((16,), eps=1e-05, elementwise_affine=True)\n",
       "            (2): GELU(approximate='none')\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import src.networks\n",
    "from physicsnemo.models.figconvnet.geometries import GridFeaturesMemoryFormat\n",
    "\n",
    "\n",
    "model = src.networks.FIGConvUNetDrivAerNet(\n",
    "    aabb_max=[2.75, 1.5, 1.0],\n",
    "    aabb_min=[-2.75, -1.5, -1.0],\n",
    "    hidden_channels=[16, 16, 16],\n",
    "    in_channels=1,\n",
    "    kernel_size=5,\n",
    "    mlp_channels=[2048, 2048],\n",
    "    neighbor_search_type=\"radius\",\n",
    "    num_down_blocks=1,\n",
    "    num_levels=2,\n",
    "    out_channels=1,\n",
    "    pooling_layers=[2],\n",
    "    pooling_type=\"max\",\n",
    "    reductions=[\"mean\"],\n",
    "    resolution_memory_format_pairs=[\n",
    "        (GridFeaturesMemoryFormat.b_xc_y_z, [5, 150, 100]),\n",
    "        (GridFeaturesMemoryFormat.b_yc_x_z, [250, 3, 100]),\n",
    "        (GridFeaturesMemoryFormat.b_zc_x_y, [250, 150, 2]),\n",
    "    ],\n",
    "    use_rel_pos_encode=True,\n",
    ")\n",
    "# Load checkpoint.\n",
    "chk = torch.load(\"/data/src/physicsnemo/models/fignet/2060187/model_00103.pth\")\n",
    "model.load_state_dict(chk[\"model\"])\n",
    "model = model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run inference\n",
    "\n",
    "The following code runs inference for the first sample in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Module physicsnemo.models.figconvnet.warp_neighbor_search load on device 'cuda:0' took 153.56 ms\n"
     ]
    }
   ],
   "source": [
    "torch.set_grad_enabled(False)\n",
    "\n",
    "sample = next(iter(dataset.train_dataloader()))\n",
    "vertices = model.data_dict_to_input(sample)\n",
    "pressure = sample[\"time_avg_pressure\"]\n",
    "normalized_pred, drag_pred = model(vertices)\n",
    "pred = dataset.normalizer.decode(normalized_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize inputs and predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "22df22d1b0f34fbd93f2123e6c4fb274",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Widget(value='<iframe src=\"http://localhost:42125/index.html?ui=P_0x7f2144ef6d10_4&reconnect=auto\" class=\"pyvi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "camera_position = [\n",
    "    (-3.0, -4.5, 5),\n",
    "    (0.5, 0, 0.6),\n",
    "    (0.1, 0.1, 0.9),\n",
    "]\n",
    "\n",
    "\n",
    "def plot_results(\n",
    "    vertices: np.ndarray,\n",
    "    pred: np.ndarray,\n",
    "    gt: np.ndarray,\n",
    "    camera_position,\n",
    "    scalar_name: str = \"p\",\n",
    "):\n",
    "    plotter = pv.Plotter(shape=(1, 3))\n",
    "\n",
    "    pc = pv.PolyData(vertices)\n",
    "    gt_name = scalar_name + \"_gt\"\n",
    "    pc[gt_name] = gt\n",
    "    pred_name = scalar_name + \"_pred\"\n",
    "    pc[pred_name] = pred\n",
    "\n",
    "    plotter.subplot(0, 0)\n",
    "    plotter.add_points(pv.PolyData(vertices))\n",
    "    plotter.camera_position = camera_position\n",
    "    plotter.camera.zoom(0.5)\n",
    "    plotter.add_text(\"Input point cloud\", position=\"upper_left\")\n",
    "\n",
    "    plotter.subplot(0, 1)\n",
    "    plotter.add_mesh(pc, scalars=pred_name, cmap=\"jet\", show_scalar_bar=False)\n",
    "    plotter.camera_position = camera_position\n",
    "    plotter.camera.zoom(0.5)\n",
    "    plotter.add_scalar_bar(title=scalar_name, vertical=True)\n",
    "    plotter.add_text(\"Predicted point cloud\", position=\"upper_left\")\n",
    "\n",
    "    plotter.subplot(0, 2)\n",
    "    plotter.add_mesh(pc, scalars=gt_name, cmap=\"jet\", show_scalar_bar=False)\n",
    "    plotter.camera_position = camera_position\n",
    "    plotter.camera.zoom(0.5)\n",
    "    plotter.add_scalar_bar(title=scalar_name + \" \", vertical=True)\n",
    "    plotter.add_text(\"GT point cloud\", position=\"upper_left\")\n",
    "\n",
    "    plotter.show()\n",
    "\n",
    "\n",
    "plot_results(\n",
    "    vertices[0].cpu().numpy(),\n",
    "    pred[0].cpu().numpy(),\n",
    "    pressure[0].cpu().numpy(),\n",
    "    camera_position,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
